#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os, re, argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def find_header_start(lines):
    """
    返回数据真正表头行的索引：
    - 优先找含 'Frame#' 的行
    - 其次找包含 X1/Y1/Z1 或 X3/Y3/Z3 的行
    - 最后退化：找到第一行看起来像“数字表”的行
    """
    for i, ln in enumerate(lines):
        if "Frame#" in ln:
            return i
    for i, ln in enumerate(lines):
        if re.search(r"\b(X1|Y1|Z1|X3|Y3|Z3)\b", ln):
            return i
    for i, ln in enumerate(lines):
        # 由数字和空白/分隔符组成的行
        if re.search(r"[0-9]", ln) and re.search(r"[,\t ]", ln):
            return i
    return 0

def read_table_any(path):
    """
    既能读标准 CSV，也能读 Motion Analysis 的 .ly：
    - 自动跳过无效头部行
    - 自动识别分隔符（优先空白，其次逗号/分号/Tab）
    """
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        lines = f.readlines()
    start = find_header_start(lines)
    tmp = path + ".tmpdata"
    with open(tmp, "w", encoding="utf-8") as f:
        f.writelines(lines[start:])

    # 优先用空白分隔（.ly 常见）
    for sep in [r"\s+", ",", ";", "\t"]:
        try:
            df = pd.read_csv(tmp, sep=sep, engine="python")
            df = df.dropna(axis=1, how="all")
            if not df.empty:
                return df.reset_index(drop=True)
        except Exception:
            pass
    raise RuntimeError(f"无法读取文件：{path}（请检查是否为文本/是否有权限）")

def find_time_column(df):
    cols = {str(c).lower(): c for c in df.columns}
    if "t" in cols: return pd.to_numeric(df[cols["t"]], errors="coerce").values, False
    for k in ["time", "timestamp", "stamp"]:
        if k in cols: return pd.to_numeric(df[cols[k]], errors="coerce").values, False
    if "sec" in cols and ("nsec" in cols or "nanosec" in cols):
        sec = pd.to_numeric(df[cols["sec"]], errors="coerce").fillna(0).values
        nkey = "nsec" if "nsec" in cols else "nanosec"
        nsec = pd.to_numeric(df[cols[nkey]], errors="coerce").fillna(0).values
        return (sec + nsec*1e-9), False
    return np.arange(len(df)), True  # 用样本索引

def pick_xyz(df, wanted):
    """大小写不敏感拿列；拿不到就兜底取前三列。"""
    m = {str(c).lower(): c for c in df.columns}
    cols = [m.get(name.lower()) for name in wanted]
    if all(c is not None for c in cols):
        return df[cols].apply(pd.to_numeric, errors="coerce").values
    if df.shape[1] >= 3:
        return df.iloc[:, :3].apply(pd.to_numeric, errors="coerce").values
    raise RuntimeError(f"未找到列 {wanted} 且列数不足 3。")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("face_file")
    ap.add_argument("arm_file")
    ap.add_argument("--face-mm", action="store_true", help="face 的 X1/Y1/Z1 单位 mm→m")
    ap.add_argument("--arm-mm", action="store_true", help="arm 的 X3/Y3/Z3 单位 mm→m")
    ap.add_argument("--out", default=".", help="输出目录")
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)

    face_df = read_table_any(args.face_file)
    arm_df  = read_table_any(args.arm_file)

    t_face, face_idx = find_time_column(face_df)
    t_arm,  arm_idx  = find_time_column(arm_df)

    face_xyz = pick_xyz(face_df, ["X1","Y1","Z1"])
    arm_xyz  = pick_xyz(arm_df,  ["X3","Y3","Z3"])

    if args.face_mm: face_xyz = face_xyz / 1000.0
    if args.arm_mm:  arm_xyz  = arm_xyz  / 1000.0

    # 图1：face
    plt.figure()
    plt.plot(t_face, face_xyz[:,0], label="X1")
    plt.plot(t_face, face_xyz[:,1], label="Y1")
    plt.plot(t_face, face_xyz[:,2], label="Z1")
    plt.xlabel("time (s)" if not face_idx else "sample")
    plt.ylabel("value")
    plt.title("face: X1/Y1/Z1")
    plt.legend(); plt.tight_layout()
    plt.savefig(os.path.join(args.out, "face_xyz.png")); plt.close()

    # 图2：arm
    plt.figure()
    plt.plot(t_arm, arm_xyz[:,0], label="X3")
    plt.plot(t_arm, arm_xyz[:,1], label="Y3")
    plt.plot(t_arm, arm_xyz[:,2], label="Z3")
    plt.xlabel("time (s)" if not arm_idx else "sample")
    plt.ylabel("value")
    plt.title("arm: X3/Y3/Z3")
    plt.legend(); plt.tight_layout()
    plt.savefig(os.path.join(args.out, "arm_xyz.png")); plt.close()

    # 图3：3D 叠加
    try:
        from mpl_toolkits.mplot3d import Axes3D  # noqa
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
        ax.plot(face_xyz[:,0], face_xyz[:,1], face_xyz[:,2], label="face (X1,Y1,Z1)")
        ax.plot(arm_xyz[:,0],  arm_xyz[:,1],  arm_xyz[:,2],  label="arm (X3,Y3,Z3)")
        ax.set_xlabel("X"); ax.set_ylabel("Y"); ax.set_zlabel("Z")
        ax.set_title("3D Trajectories")
        ax.legend(); plt.tight_layout()
        plt.savefig(os.path.join(args.out, "face_arm_3d.png")); plt.close()
    except Exception as e:
        print("3D 绘图失败（可忽略）：", e)

    print("完成：face_xyz.png, arm_xyz.png, face_arm_3d.png")

if __name__ == "__main__":
    main()
